package com.zhuqian.myapplication

import android.content.Context
import android.graphics.Bitmap
import android.graphics.Rect
import org.pytorch.IValue
import org.pytorch.LiteModuleLoader
import org.pytorch.torchvision.TensorImageUtils
import java.io.File
import java.io.FileOutputStream

data class Result(var classIndex: Int, var score: Float, var rect: Rect);

class Detector(
    ctx: Context,
    path: String,
    val img_size: Int,
    val type: String,
    val class_num: Int = 1,
    var iou: Float = 0.3f,
    var conf: Float = 0.3f,
    var nms_limit: Int = 500,
) {

    val NO_MEAN_RGB = floatArrayOf(0.0f, 0.0f, 0.0f)
    val NO_STD_RGB = floatArrayOf(1.0f, 1.0f, 1.0f)

    val model = LiteModuleLoader.load(assetFilePath(ctx, path))
    fun assetFilePath(context: Context, assetName: String): String {
        val file = File(context.filesDir, assetName)
        if (file.exists() && file.length() > 0) {
            return file.absolutePath
        }
        context.assets.open(assetName).use { `is` ->
            FileOutputStream(file).use { os ->
                val buffer = ByteArray(4 * 1024)
                var read: Int
                while (`is`.read(buffer).also { read = it } != -1) {
                    os.write(buffer, 0, read)
                }
                os.flush()
            }
            return file.absolutePath
        }
    }


    fun detect_bitmap(bitmap: Bitmap): ArrayList<Result> {
        // 缩放比例
        val xcaleX: Float = bitmap.width.toFloat() / img_size
        val xcaleY: Float = bitmap.height.toFloat() / img_size
        // 缩放图片
        val resizedBitmap =
            Bitmap.createScaledBitmap(bitmap, img_size, img_size, true)
        // 转换图片
        val inputTensor = TensorImageUtils.bitmapToFloat32Tensor(
            resizedBitmap,
            NO_MEAN_RGB,
            NO_STD_RGB
        )
        val input = IValue.from(inputTensor)
        // 模型预测
        val outputs = model.forward(input)
        return Processor.parse_output(
            outputs,
            class_num,
            xcaleX,
            xcaleY,
            nms_limit,
            iou,
            conf,
            type
        )
    }
}

